package Redis

/*
	基于Redis实现的分布式锁，保障在分布式环境下同一业务只能被执行一次
	Acquire(5) 在设定时间内抢占到锁，每隔一段时间会自动刷新锁直至调用 Release()
	Exp:

	password := "12345"
	Pool := redis.Pool{
		MaxIdle: 2,
		MaxActive: 2,
		IdleTimeout: 6 * time.Second,
		// Dial or DialContext must be set. When both are set, DialContext takes precedence over Dial.
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", "127.0.0.1:6379")
			if err != nil {
				return c, errors.New(fmt.Sprintf("redis connection error: %s", err))
			}
			if password != "" {
				if _, authErr := c.Do("AUTH", password); authErr != nil {
					return c, fmt.Errorf("redis auth password error: %s", authErr)
				}
			}
			//选择分区
			c.Do("SELECT", 0)
			return c, nil
		},
	}
	var Pools []*redis.Pool

	Pools = append(Pools, &Pool)
	//导入InitMLock，并全局初始化MLock
	locker, _ := InitMLock(&Pools, 30) // 可以传入多个redis连接池，会根据键的hash值求余选择对应的redis连接池
	locker.Start()

	// 导入Mlock
	mlock := Mlock.GetLock("test")
	if mlock.Acquire(5){
		fmt.Println("抢到锁了")
		// do something
	}else{
		fmt.Println("没抢到锁")
	}
	locked := mlock.Locked()
	fmt.Println(locked)
	locked.Release()
*/

import (
	"errors"
	"fmt"
	"gitee.com/xixigroup/sweet-tools/base/Utils"
	"github.com/gomodule/redigo/redis"
	"runtime"
	"time"
)

const (
	RENEW_SCRIPT = `
		if redis.call("get",KEYS[1]) == ARGV[1] and redis.call("ttl",KEYS[1]) > 0 then
			return redis.call("expire",KEYS[1],ARGV[2])
		else
			return 0
		end
	`
	UNLOCK_SCRIPT = `
		if redis.call("get",KEYS[1]) == ARGV[1] then
			return redis.call("del",KEYS[1])
		else
			return 0
		end
	`
)

func InitMLock(pools *[]*redis.Pool, expire int) (*MLock, error) {
	// 创建锁
	if expire < 30 {
		return nil, errors.New("expire cannot less then 30")
	}

	Mlock = &MLock{pools: pools}
	Mlock.lockTag = "MLOCK_PROCESS_"
	Mlock.lockVal = 1
	Mlock.expire = expire
	Mlock.renewInterval = expire / 3
	Mlock.locks = Utils.NewSafeMap()

	// 析构lock前 清理收尾工作
	runtime.SetFinalizer(Mlock, recycle)
	go Mlock.run()
	return Mlock, nil
}

type MLock struct {
	pools         *[]*redis.Pool
	lockTag       string
	lockVal       interface{}
	expire        int
	locks         *Utils.SafeMap
	renewInterval int
	running       bool
}

var Mlock *MLock

type Lock struct {
	key    string
	val    interface{}
	expire int
	locked bool
	marks  uint8
	hash   uint32
}

func (locker *MLock) GetLock(key string) *Lock {
	key = locker.lockTag + key

	lock := Lock{}
	lock.key = key
	lock.val = locker.lockVal
	lock.expire = locker.expire
	lock.hash = Utils.HashMath(lock.key)

	locks := locker.locks
	locks.Set(key, &lock)
	return &lock
}

func (locker *MLock) Start() {
	locker.running = true
}

func (locker *MLock) Stop() {
	locker.running = false
}

func recycle(locker *MLock) {
	// 资源回收
	locker.Close()
}

func (locker *MLock) Close() {
	// 资源回收
	locker.running = false
	locker.locks = Utils.NewSafeMap()
	locker.pools = nil
}

func (locker *MLock) run() {

	for {
		if locker.running {
			go func() {

				for _, lock := range locker.locks.Values() {

					switch l := lock.(type) {
					case *Lock:
						if l.locked {
							Utils.ProtectRunBoolFunc(l.renew)
						}
					default:
						continue
					}
				}
			}()
		}

		time.Sleep(time.Second * time.Duration(locker.renewInterval))
	}
}

func (lock *Lock) getConn() *redis.Conn {
	conn := (*Mlock.pools)[lock.chooseRedisPool()].Get()
	return &conn
}

func (lock *Lock) chooseRedisPool() int {

	return int(lock.hash) % len(*Mlock.pools)
}

func (lock *Lock) Exist() bool {
	// 锁是否存在
	conn := lock.getConn()
	defer (*conn).Close()
	res, _ := (*conn).Do("EXISTS", lock.key)
	switch r := res.(type) {
	case int64:
		if r == int64(1) {
			return true
		}
	}
	return false
}

func (lock *Lock) Release() error {
	// 释放锁
	var err error
	if lock.locked {
		lock.locked = false
		conn := lock.getConn()
		defer (*conn).Close()

		lua := redis.NewScript(1, UNLOCK_SCRIPT) // 1 决定Do的前几个为key

		ok, _ := lua.Do(*conn, lock.key, lock.val)

		switch o := ok.(type) {
		case int64:
			if o == int64(1) {
				err = nil
			}
		default:
			err = errors.New(fmt.Sprintf("MLock Release Error key: %s err:%v", lock.key, err))
		}
	}
	return err
}

func (lock *Lock) Locked() bool {
	// 检查持有锁状态
	return lock.locked
}

func (lock *Lock) renew() bool {
	// 刷新锁
	if lock.locked {
		conn := lock.getConn()
		defer (*conn).Close()
		lua := redis.NewScript(1, RENEW_SCRIPT) // 1 决定Do的前几个为key

		ok, err := lua.Do(*conn, lock.key, lock.val, lock.expire)

		if err != nil {
			lock.locked = false
		}
		switch flag := ok.(type) {
		case int64:
			if flag == int64(1) {
				lock.locked = true
			} else {
				lock.locked = false
			}
		default:
			lock.locked = false
		}
	}

	return lock.locked
}

func (lock *Lock) Acquire(timeout uint) bool {
	// 抢占锁，在超时时间内去抢占
	defer Utils.TracebackFormatExec(fmt.Sprintf("lock[%s] acquire error", lock.key))

	if lock.locked {
		return lock.renew()
	} else {
		return lock.acquire(timeout)
	}
}

func (lock *Lock) acquire(timeout uint) bool {

	conn := lock.getConn()
	defer (*conn).Close()

	t := uint(time.Now().Unix()) + timeout
	for {
		res, _ := (*conn).Do("SET", lock.key, lock.val, "EX", lock.expire, "NX")
		switch r := res.(type) {
		case string:

			if string(r) == "OK" {
				lock.locked = true
				return lock.locked
			}
		}
		if timeout > 0 && t > uint(time.Now().Unix()) {
			time.Sleep(time.Second)
		} else {
			break
		}
	}
	return lock.locked
}
